/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.util;



import java.io.ByteArrayOutputStream;

import java.io.IOException;



/**

 * {@link ByteArrayOutputStream} special cased to treat writes of a single byte-array specially.

 * When calling {@link #toByteArray()} after writing only one {@code byte[]} using {@link

 * #writeAndOwn(byte[])}, it will return that array directly.

 */

public class ExposedByteArrayOutputStream extends ByteArrayOutputStream {



  private byte[] swappedBuffer;



  /**

   * If true, this stream doesn't allow direct access to the passed in byte-array. It behaves just

   * like a normal {@link ByteArrayOutputStream}.

   *

   * <p>It is set to true after any write operations other than the first call to {@link

   * #writeAndOwn(byte[])}.

   */

  private boolean isFallback = false;



  /** Fall back to the behavior of a normal {@link ByteArrayOutputStream}. */

  private void fallback() {

    isFallback = true;

    if (swappedBuffer != null) {

      // swappedBuffer != null means buf is actually provided by the caller of writeAndOwn(),

      // while swappedBuffer is the original buffer.

      // Recover the buffer and copy the bytes from buf.

      byte[] tempBuffer = buf;

      count = 0;

      buf = swappedBuffer;

      super.write(tempBuffer, 0, tempBuffer.length);

      swappedBuffer = null;

    }

  }



  /**

   * Write {@code b} to the stream and take the ownership of {@code b}. If the stream is empty,

   * {@code b} itself will be used as the content of the stream and no content copy will be

   * involved.

   *

   * <p><i>Note: After passing any byte array to this method, it must not be modified again.</i>

   */

  public synchronized void writeAndOwn(byte[] b) throws IOException {

    if (b.length == 0) {

      return;

    }

    if (count == 0) {

      // Optimized first-time whole write.

      // The original buffer will be swapped to swappedBuffer, while the input b is used as buf.

      swappedBuffer = buf;

      buf = b;

      count = b.length;

    } else {

      fallback();

      super.write(b);

    }

  }



  @Override

  public synchronized void write(byte[] b, int off, int len) {

    fallback();

    super.write(b, off, len);

  }



  @Override

  public synchronized void write(int b) {

    fallback();

    super.write(b);

  }



  @Override

  public synchronized byte[] toByteArray() {

    // Note: count == buf.length is not a correct criteria to "return buf;", because the internal

    // buf may be reused after reset().

    if (!isFallback && count > 0) {

      return buf;

    } else {

      return super.toByteArray();

    }

  }



  @Override

  public synchronized void reset() {

    if (count == 0) {

      return;

    }

    count = 0;

    if (isFallback) {

      isFallback = false;

    } else {

      buf = swappedBuffer;

      swappedBuffer = null;

    }

  }

}